from datetime import datetime, timedelta
from os.path import exists

from Crypto import Hash
from Crypto.Cipher import PKCS1_v1_5 as PKCS1_v1_5_Cipher
from Crypto.PublicKey import RSA
from Crypto.Signature import PKCS1_v1_5 as PKCS1_v1_5_Signature
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateNumbers
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicNumbers
from cryptography.hazmat.primitives.asymmetric.rsa import generate_private_key
from cryptography.hazmat.primitives.asymmetric.rsa import rsa_recover_prime_factors
from cryptography.x509 import Certificate
from cryptography.x509 import CertificateBuilder
from cryptography.x509 import Name
from cryptography.x509 import NameAttribute
from cryptography.x509 import load_pem_x509_certificate
from cryptography.x509 import random_serial_number
from cryptography.x509.oid import NameOID


class RSAClient:
    def __init__(self, key: RSAPrivateKey):
        self.key = key  # 使用openssl加速
        private_numbers = key.private_numbers()
        self.private_key = RSA.construct(
            (
                private_numbers.public_numbers.n,
                private_numbers.public_numbers.e,
                private_numbers.d,
                private_numbers.p,
                private_numbers.q,
            )
        )
        self.public_key = self.private_key.publickey()

    @classmethod
    def generate(cls, nbits=4096, public_exponent=65537) -> 'RSAClient':
        key = generate_private_key(public_exponent, nbits)
        return cls(key)

    @classmethod
    def create(
            cls,
            *,
            e=65537, d: int | None = None,
            n: int | None = None,
            p: int | None = None,
            q: int | None = None
    ) -> 'RSAClient':
        if not p or not q:
            p, q = rsa_recover_prime_factors(n, e, d)
        if not d or not n:
            n = p * q
            d = pow(e, -1, (p - 1) * (q - 1))
        key = RSAPrivateNumbers(p, q, d, d % (p - 1), d % (q - 1), pow(q, -1, p), RSAPublicNumbers(e, n)).private_key()
        return cls(key)

    def dump(self, key_path="ca.key"):
        self.dump_private_key(key_path)

    def dump_private_key(self, path="ca.key"):
        with open(path, "wb") as key_file:
            key_file.write(self.private_key.export_key())

    @classmethod
    def load(cls, key_path="ca.key", password=None) -> 'RSAClient':
        return cls(RSAClient.load_private_key(key_path, password))

    @staticmethod
    def load_private_key(path="ca.key", password=None):
        if not exists(path):
            raise Exception(f"{path} not found")
        with open(path, "rb") as key_file:
            return serialization.load_pem_private_key(key_file.read(), password)

    @staticmethod
    def generate_cert(key: RSAPrivateKey, subject_name, issuer_name, days=36500) -> Certificate:
        today = datetime.today()
        one_day = timedelta(days=1)
        time_range = timedelta(days=days)
        start_time = today - one_day
        end_time = today + time_range

        builder = CertificateBuilder()
        builder = builder.subject_name(Name([
            NameAttribute(NameOID.COMMON_NAME, subject_name),
        ]))
        builder = builder.issuer_name(Name([
            NameAttribute(NameOID.COMMON_NAME, issuer_name),
        ]))
        builder = builder.not_valid_before(start_time)
        builder = builder.not_valid_after(end_time)
        builder = builder.serial_number(random_serial_number())
        builder = builder.public_key(key.public_key())

        return builder.sign(
            private_key=key, algorithm=hashes.SHA256()
        )

    @staticmethod
    def dump_cert(certificate: Certificate, path="ca.crt"):
        with open(path, "wb") as cert_file:
            cert_file.write(RSAClient.export_cert(certificate))

    @staticmethod
    def export_cert(certificate):
        return certificate.public_bytes(serialization.Encoding.PEM)

    @staticmethod
    def load_cert(path="ca.crt"):
        if not exists(path):
            raise Exception(f"{path} not found")
        with open(path, "rb") as cert_file:
            return load_pem_x509_certificate(cert_file.read())

    # 常规方法
    def encrypt(self, message):
        cipher = PKCS1_v1_5_Cipher.new(self.public_key)
        return cipher.encrypt(message)

    def decrypt(self, message):
        cipher = PKCS1_v1_5_Cipher.new(self.private_key)
        return cipher.decrypt(message, None)

    def sign(self, message, hash_algorithm='SHA1'):
        hash_algorithm = getattr(Hash, hash_algorithm)
        hashed_message = hash_algorithm.new(message)
        signer = PKCS1_v1_5_Signature.new(self.private_key)
        signature = signer.sign(hashed_message)
        return signature

    def verify(self, message, signature, hash_algorithm='SHA1'):
        hash_algorithm = getattr(Hash, hash_algorithm)
        hashed_message = hash_algorithm.new(message)
        verifier = PKCS1_v1_5_Signature.new(self.public_key)
        return verifier.verify(hashed_message, signature)


if __name__ == '__main__':
    obj = RSAClient.generate()
    msg = b'asd'
    assert obj.verify(msg, obj.sign(msg))
    assert obj.decrypt(obj.encrypt(msg)) == msg
    obj.dump()
    pass
